#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AESAlgorithm的泄漏模型
定义AES各个操作的Power泄漏特征

架构Description:
1. Intermediate valueFunction (Intermediate Functions): Calculate加密过程中的Intermediate value
2. 泄漏模型 (Leakage Models): 将Intermediate valueConvertFor泄漏值（HW, HDetc.）
3. 目标操作 (Target Operations): 可Attack的AES操作阶段
"""

import numpy as np
from typing import Dict, Any, List, Callable


class AESLeakageModels:
    """AES泄漏模型集合"""
    
    # AES S-box
    AES_SBOX = np.array([
        0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
        0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
        0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
        0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
        0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
        0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
        0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
        0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
        0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
        0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
        0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
        0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
        0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
        0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
        0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
        0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
    ], dtype=np.uint8)
    
    # AES逆S-box
    AES_INV_SBOX = np.array([
        0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb,
        0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f, 0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb,
        0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e,
        0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25,
        0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d, 0x65, 0xb6, 0x92,
        0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda, 0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84,
        0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06,
        0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b,
        0x3a, 0x91, 0x11, 0x41, 0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73,
        0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e,
        0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b,
        0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd, 0x5a, 0xf4,
        0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f,
        0x60, 0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef,
        0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5, 0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61,
        0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d
    ], dtype=np.uint8)
    
    @staticmethod
    def hamming_weight(value: int) -> int:
        """Calculate汉明权重"""
        return bin(value).count('1')
    
    @staticmethod
    def hamming_distance(a: int, b: int) -> int:
        """CalculateHamming distance"""
        return AESLeakageModels.hamming_weight(a ^ b)
    
    @classmethod
    def sbox_output_model(cls, plaintext_byte: int, key_byte: int) -> int:
        """
        AES S-boxOutput泄漏模型
        
        Args:
            plaintext_byte: Plaintext字节
            key_byte: Key字节
            
        Returns:
            S-boxOutput值
        """
        return cls.AES_SBOX[plaintext_byte ^ key_byte]
    
    @classmethod
    def sbox_output_hamming_weight(cls, plaintext_byte: int, key_byte: int) -> int:
        """
        AES S-boxOutput汉明权重泄漏模型
        
        Args:
            plaintext_byte: Plaintext字节
            key_byte: Key字节
            
        Returns:
            S-boxOutput的汉明权重
        """
        sbox_output = cls.sbox_output_model(plaintext_byte, key_byte)
        return cls.hamming_weight(sbox_output)
    
    @classmethod
    def sbox_input_xor_model(cls, plaintext_byte: int, key_byte: int) -> int:
        """
        AES S-boxInputXOR泄漏模型
        
        Args:
            plaintext_byte: Plaintext字节
            key_byte: Key字节
            
        Returns:
            S-boxInput值（PlaintextXORKey）
        """
        return plaintext_byte ^ key_byte
    
    @classmethod
    def sbox_input_hamming_weight(cls, plaintext_byte: int, key_byte: int) -> int:
        """
        AES S-boxInput汉明权重泄漏模型
        
        Args:
            plaintext_byte: Plaintext字节
            key_byte: Key字节
            
        Returns:
            S-boxInput的汉明权重
        """
        sbox_input = cls.sbox_input_xor_model(plaintext_byte, key_byte)
        return cls.hamming_weight(sbox_input)
    
    @classmethod
    def round_key_xor_model(cls, state_byte: int, round_key_byte: int) -> int:
        """
        轮KeyXOR泄漏模型
        
        Args:
            state_byte: Status字节
            round_key_byte: 轮Key字节
            
        Returns:
            XORResult
        """
        return state_byte ^ round_key_byte
    
    @classmethod
    def shift_rows_model(cls, state: List[List[int]]) -> List[List[int]]:
        """
        ShiftRows操作泄漏模型
        
        Args:
            state: 4x4StatusMatrix（按行存储）
            
        Returns:
            ShiftRows后的Status
        """
        if len(state) != 4 or any(len(row) != 4 for row in state):
            raise ValueError("Status必须Yes4x4Matrix")
        
        # Create新Status
        result = [row[:] for row in state]
        
        # 第0行Not移位
        # 第1行左移1字节
        result[1] = state[1][1:] + state[1][:1]
        # 第2行左移2字节
        result[2] = state[2][2:] + state[2][:2]
        # 第3行左移3字节
        result[3] = state[3][3:] + state[3][:3]
        
        return result
    
    @classmethod
    def sub_bytes_model(cls, state: List[List[int]]) -> List[List[int]]:
        """
        SubBytes操作泄漏模型（对整个StatusMatrix）
        
        Args:
            state: 4x4StatusMatrix
            
        Returns:
            SubBytes后的Status
        """
        result = []
        for row in state:
            result.append([cls.AES_SBOX[byte] for byte in row])
        return result
    
    @classmethod
    def mixcolumns_model(cls, column: List[int]) -> List[int]:
        """
        MixColumns操作泄漏模型（单列）
        
        Args:
            column: 4字节列
            
        Returns:
            MixColumns后的列
        """
        # 简化的MixColumns模型
        # 实际实现会更复杂，这里ProvideBasic框架
        def gf_mult(a: int, b: int) -> int:
            """GF(2^8)乘法"""
            result = 0
            for _ in range(8):
                if b & 1:
                    result ^= a
                a <<= 1
                if a & 0x100:
                    a ^= 0x11b
                b >>= 1
            return result & 0xff
        
        # MixColumnsMatrix
        mix_matrix = [
            [2, 3, 1, 1],
            [1, 2, 3, 1],
            [1, 1, 2, 3],
            [3, 1, 1, 2]
        ]
        
        result = []
        for i in range(4):
            val = 0
            for j in range(4):
                val ^= gf_mult(mix_matrix[i][j], column[j])
            result.append(val)
        
        return result
    
    @classmethod
    def mixcolumns_state_model(cls, state: List[List[int]]) -> List[List[int]]:
        """
        MixColumns操作泄漏模型（整个Status）
        
        Args:
            state: 4x4StatusMatrix（按行存储）
            
        Returns:
            MixColumns后的Status
        """
        # 转置：From行存储转For列存储
        columns = [[state[row][col] for row in range(4)] for col in range(4)]
        
        # 对每列应用MixColumns
        mixed_columns = [cls.mixcolumns_model(col) for col in columns]
        
        # 转置回行存储
        result = [[mixed_columns[col][row] for col in range(4)] for row in range(4)]
        
        return result
    
    @classmethod
    def key_schedule_model(cls, key: bytes, round_num: int) -> bytes:
        """
        AESKey调度泄漏模型（AES-128标准Algorithm）
        
        实现符合FIPS 197标准的AESKey扩展Algorithm
        
        Args:
            key: 主Key（16字节，AES-128）
            round_num: 目标轮数（1-10）
            
        Returns:
            指定轮的轮Key（16字节）
            
        AlgorithmDescription:
        1. 将16字节Key视For4个字（Each4字节）
        2. 每轮扩展Generate4个新字
        3. 第i轮Key = 第i-1轮Key的Function
        4. 关键操作: RotWord, SubWord, Rcon
        """
        if len(key) != 16:
            raise ValueError("Key必须Yes16字节（AES-128）")
        
        if not (1 <= round_num <= 10):
            raise ValueError("轮数必须In1-10之间")
        
        # Rcon轮常数（只Require前10个）
        rcon = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36]
        
        # InitializeKey调度Array（11轮，每轮16字节）
        w = [0] * (4 * (round_num + 1))
        
        # 第0轮：复制原始Key
        for i in range(4):
            w[i] = int.from_bytes(key[4*i:4*i+4], 'big')
        
        # 扩展Key
        for i in range(4, 4 * (round_num + 1)):
            temp = w[i - 1]
            
            if i % 4 == 0:
                # 每4个字Execute特殊操作
                # 1. RotWord: Loop左移1字节
                temp = ((temp << 8) | (temp >> 24)) & 0xffffffff
                
                # 2. SubWord: 对Each字节应用S-box
                temp_bytes = [
                    int(cls.AES_SBOX[(temp >> 24) & 0xff]),
                    int(cls.AES_SBOX[(temp >> 16) & 0xff]),
                    int(cls.AES_SBOX[(temp >> 8) & 0xff]),
                    int(cls.AES_SBOX[temp & 0xff])
                ]
                temp = (temp_bytes[0] << 24) | (temp_bytes[1] << 16) | (temp_bytes[2] << 8) | temp_bytes[3]
                
                # 3. 与Rcon异Or（只Has最高字节）
                rcon_word = (rcon[(i // 4) - 1] << 24) & 0xffffffff
                temp = (temp ^ rcon_word) & 0xffffffff
            
            w[i] = (w[i - 4] ^ temp) & 0xffffffff
        
        # Extract指定轮的Key（Last4个字）
        round_key = bytearray(16)
        for i in range(4):
            word = w[4 * round_num + i]
            round_key[4*i:4*i+4] = word.to_bytes(4, 'big')
        
        return bytes(round_key)
    
    @classmethod
    def key_schedule_intermediate_model(cls, prev_key: bytes, round_num: int) -> bytes:
        """
        Key调度Intermediate value泄漏模型（单步扩展）
        
        Args:
            prev_key: 前一轮的Key（16字节）
            round_num: Current轮数（1-10）
            
        Returns:
            Current轮的Key（16字节）
        """
        # UseComplete的Key调度，但只ReturnLast一轮
        # 这里For了Attack建模，我们RequireFrom主KeyStart推导
        # 实际Attack中，可能只针对某个特定的Key扩展Step
        
        rcon = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36]
        
        # 将prev_key分For4个字
        w = [int.from_bytes(prev_key[4*i:4*i+4], 'big') for i in range(4)]
        
        # 扩展4个新字
        w_new = []
        for i in range(4):
            if i == 0:
                # First个字: Require特殊Process
                temp = w[3]  # Last一个字
                
                # RotWord
                temp = ((temp << 8) | (temp >> 24)) & 0xffffffff
                
                # SubWord
                temp_bytes = [
                    cls.AES_SBOX[(temp >> 24) & 0xff],
                    cls.AES_SBOX[(temp >> 16) & 0xff],
                    cls.AES_SBOX[(temp >> 8) & 0xff],
                    cls.AES_SBOX[temp & 0xff]
                ]
                temp = (temp_bytes[0] << 24) | (temp_bytes[1] << 16) | (temp_bytes[2] << 8) | temp_bytes[3]
                
                # Rcon
                temp ^= (rcon[round_num - 1] << 24)
                
                w_new.append(w[0] ^ temp)
            else:
                w_new.append(w[i] ^ w_new[i - 1])
        
        # 组合成16字节
        result = bytearray(16)
        for i in range(4):
            result[4*i:4*i+4] = w_new[i].to_bytes(4, 'big')
        
        return bytes(result)
    
    # ===================================================================
    # 新架构：分离三个核心概念
    # ===================================================================
    
    @classmethod
    def get_intermediate_functions(cls) -> Dict[str, Callable]:
        """
        获取Intermediate valueCalculateFunction（Calculate什么Intermediate value）
        这些Function定义如何FromPlaintextAndKeyCalculate加密过程中的Intermediate value
        
        Returns:
            Intermediate valueCalculateFunctionDictionary {名称: Function}
        """
        return {
            # S-box相关
            'sbox_output': cls.sbox_output_model,         # S-boxOutput值
            'sbox_input': cls.sbox_input_xor_model,       # S-boxInput值（P⊕K）
            
            # 轮操作相关
            'round_key_xor': cls.round_key_xor_model,     # 轮Key异Or
            'shift_rows': cls.shift_rows_model,           # ShiftRows操作
            'sub_bytes': cls.sub_bytes_model,             # SubBytes操作（整个Status）
            'mixcolumns': cls.mixcolumns_model,           # MixColumns操作（单列）
            'mixcolumns_state': cls.mixcolumns_state_model,  # MixColumns操作（整个Status）
            'key_schedule': cls.key_schedule_model,       # Key调度
            'key_schedule_step': cls.key_schedule_intermediate_model,  # 单步Key扩展
        }
    
    @classmethod
    def get_leakage_models(cls) -> Dict[str, Callable]:
        """
        获取泄漏表示模型（如何泄漏）
        这些Function定义如何将Intermediate valueConvertFor可观察的泄漏值
        
        Returns:
            泄漏模型Dictionary {名称: Function}
        """
        return {
            'value': lambda x: x,                         # 原始值（No变换）
            'hamming_weight': cls.hamming_weight,         # 汉明权重
            'hamming_distance': cls.hamming_distance,     # Hamming distance
            'bit_0': lambda x: x & 1,                     # 最低位
            'bit_7': lambda x: (x >> 7) & 1,             # 最高位
        }
    
    @classmethod
    def get_target_operations(cls) -> List[str]:
        """
        获取可Attack的目标操作（Attack什么）
        这些YesAES加密过程中可以作ForAttack目标的操作阶段
        
        Returns:
            目标操作List
        """
        return [
            'first_round_sbox',      # First轮S-box操作
            'last_round_sbox',       # Last一轮S-box操作
            'round_key_addition',    # 轮Key加操作
            'mixcolumns',            # MixColumns操作
            'key_schedule',          # Key调度操作
            'sub_bytes',             # SubBytes操作
            'shift_rows',            # ShiftRows操作
        ]
    
    @classmethod
    def get_operation_to_intermediate_mapping(cls) -> Dict[str, str]:
        """
        获取目标操作ToIntermediate valueFunction的Mapping
        定义Each目标操作应该Use哪个Intermediate valueCalculateFunction
        
        Returns:
            MappingDictionary {目标操作: Intermediate valueFunction名}
        """
        return {
            'first_round_sbox': 'sbox_output',
            'last_round_sbox': 'sbox_output',
            'round_key_addition': 'round_key_xor',
            'mixcolumns': 'mixcolumns',
            'key_schedule': 'key_schedule',
            'sub_bytes': 'sbox_output',
            'shift_rows': 'sbox_output',
        }
    
